# 使用堆优化prim算法

# 1寻找最小值使用堆优化
# 2邻接表优化无向图的存储

# Prim算法的时间复杂度为O(N2)，如果借助堆，每次选边的时间复杂度是O(logM),然后用邻接表来存储图的话，整个时间复杂度会降到O(MlogN)

# 1、数组dis用来记录生成树到各个顶点的距离
# 2、数组h是一个最小堆，堆里面存储的是顶点编号。（请注意，这里并不是按照顶点编号的大小来建立最小堆的，而是按照顶点在数组dis中所对应的值来建立这个最小堆的）
# 3、pos数组来记录每个顶点在最小堆中的位置
dis = [0 for i in range(0, 7)]
# book数组用来记录哪些顶点已经放入生成树中
book = [0 for i in range(0, 7)]
# h用来保存堆，pos用来存储每个顶点在堆中的位置,size为堆的大小
h = [0 for i in range(0, 7)]
pos = [0 for i in range(0, 7)]
size = 7
n = 6
m = 9


# 交换数组的位置
def swap(x, y):
    temp = h[x]
    h[x] = h[y]
    h[y] = temp

    # 同步更新pos
    t = pos[h[x]]
    pos[h[x]] = pos[h[y]]
    pos[h[y]] = t


# 向下调整


def sift_down(i: int):
    t = 0  # 记录最小节点
    flag = 0  # flag用来标记是否需要向下调整
    # 当i节点有子节点（至少有子节点的情况）并且需要继续调整的时候，循环执行
    while i * 2 <= n and flag == 0:
        if dis[h[i]] > dis[h[i * 2]]:
            t = i * 2
        else:
            t = i
        # 如果有右节点
        if i * 2 + 1 <= n:
            if dis[h[t]] > dis[h[i * 2 + 1]]:
                t = i * 2 + 1
        # 如果发现最小的结点编号不是自己，说明子结点中有比父结点更小的
        if t != i:
            swap(t, i)
            i = t  # 更新i为刚才与它交换的儿子结点的编号、便于接下来继续向下调整
        else:
            flag = 1  # 否则说明当前的父结点已经比两个子结点都要小了，不需委再进行调整了


# 向上调整
def siftup(i: int):  # 专入一个霭要向上调整的结点编号i
    flag = 0
    if i == 1:
        return
    while i != 1 and flag == 0:
        if dis[h[i]] < dis[h[i // 2]]:
            swap(i, i // 2)
        else:
            flag = 1
        i = i // 2  # 当前节点的父节点的编号


# 从堆顶取一个元素
def pop():
    global size
    t = h[1]
    pos[t] = 0
    h[1] = h[size]
    pos[h[1]] = 1
    size -= 1
    sift_down(1)
    return t


if __name__ == '__main__':
    a = [[0, 0, 0], [2, 4, 11], [3, 5, 13], [4, 6, 3], [5, 6, 4], [2, 3, 6], [4, 5, 7], [1, 2, 1], [3, 4, 9], [1, 3, 2]]
    inf = 999
    count = 0
    s = 0
    # 邻接表存储图
    # u、v,w和next的数组大小要根据实际情况来设置，此图是无向图，要比2*m的最大值要大1
    u = [0 for i in range(0, 20)]
    v = [0 for i in range(0, 20)]
    w = [0 for i in range(0, 20)]
    first = [0 for i in range(0, n + 1)]
    next_ = [0 for i in range(0, 20)]
    # 初始化
    for i in range(1, len(a)):
        u[i] = a[i][0]
        v[i] = a[i][1]
        w[i] = a[i][2]

    # 这里是无向图,所以需要将所有的边再反向存储一次
    for i in range(m + 1, 2 * m + 1):
        u[i] = v[i - m]
        v[i] = u[i - m]
        w[i] = w[i - m]

    # 读入边
    for i in range(1, n + 1):
        first[i] = -1
    for i in range(1, 2 * m + 1):
        next_[i] = first[u[i]]
        first[u[i]] = i

    # print('遍历每个顶点的边')
    # for i in range(1, n + 1):
    #     k = first[i]
    #     while k != -1:
    #         print('%d,%d,%d' % (u[k], v[k], w[k]))
    #         k = next_[k]

    # prim 算法核心部分

    # 将1号顶点加入生成树
    book[1] = 1
    count += 1
    # 初始化dis数组,这里是1号顶点到其余各个顶点的初始距离
    dis[1] = 0
    for i in range(2, n + 1):
        dis[i] = inf

    k = first[1]
    while k != -1:
        dis[v[k]] = w[k]
        k = next_[k]

    # 初始化堆
    size = n
    for i in range(1, size + 1):
        h[i] = i
        pos[i] = i
    # 先弹出一个堆顶元素。因为此时堆顶是1号顶点
    for i in range(size // 2, 0, -1):
        print(i)
        sift_down(i)
    pop()

    while count < n:
        j = pop()
        book[j] = 1
        count += 1
        s += dis[j]
        # 扫描当前顶点j所有的边，再以j为中间结点。进行松弛
        k = first[j]
        while k != -1:
            if book[v[k]] == 0 and dis[v[k]] > w[k]:
                dis[v[k]] = w[k]  # 更新距离
                siftup(pos[v[k]])  # 对该顶点在堆中向上调整
                # 提示: pos[v[k】]存储的是顶点v[k】在堆中的位置
            k = next_[k]

    # 打印结果
    print('%d' % s)
